import os
import json
import numpy as np
import matplotlib.pyplot as plt


def analyze_kfold_results(work_dir_pattern, folds=5):
    results = []

    for fold in range(folds):
        work_dir = work_dir_pattern.format(fold)
        result_file = os.path.join(work_dir, 'best_coco_bbox_mAP.json')

        if os.path.exists(result_file):
            with open(result_file, 'r') as f:
                data = json.load(f)

            # 提取关键指标
            results.append({
                'fold': fold,
                'mAP': data['coco/bbox_mAP'],
                'mAP_50': data['coco/bbox_mAP_50'],
                'mAP_75': data['coco/bbox_mAP_75']
            })

    # 计算平均值和标准差
    metrics = ['mAP', 'mAP_50', 'mAP_75']
    averages = {}
    stds = {}

    for metric in metrics:
        values = [r[metric] for r in results]
        averages[metric] = np.mean(values)
        stds[metric] = np.std(values)

    # 打印结果
    print("\n===== 5折交叉验证结果 =====")
    for fold, result in enumerate(results):
        print(f"Fold {fold}: mAP={result['mAP']:.4f}, mAP_50={result['mAP_50']:.4f}, mAP_75={result['mAP_75']:.4f}")

    print("\n===== 平均结果 =====")
    for metric in metrics:
        print(f"Average {metric}: {averages[metric]:.4f} ± {stds[metric]:.4f}")

    # 可视化结果
    fold_nums = [r['fold'] for r in results]

    plt.figure(figsize=(10, 6))
    plt.bar(fold_nums, [r['mAP'] for r in results], alpha=0.7, label='mAP')
    plt.bar([x + 0.3 for x in fold_nums], [r['mAP_50'] for r in results], alpha=0.7, label='mAP_50')
    plt.bar([x + 0.6 for x in fold_nums], [r['mAP_75'] for r in results], alpha=0.7, label='mAP_75')

    plt.axhline(y=averages['mAP'], color='r', linestyle='-', alpha=0.3)
    plt.axhline(y=averages['mAP_50'], color='g', linestyle='-', alpha=0.3)
    plt.axhline(y=averages['mAP_75'], color='b', linestyle='-', alpha=0.3)

    plt.xlabel('Fold')
    plt.ylabel('Metric Value')
    plt.title('K-Fold Cross Validation Results')
    plt.xticks([r + 0.3 for r in fold_nums], [f'Fold {i}' for i in fold_nums])
    plt.legend()
    plt.tight_layout()
    plt.savefig('kfold_results.png')
    plt.show()


if __name__ == '__main__':
    work_dir_pattern = './work_dirs/diffusiondet_r50_5kfold_fpn_microalgea_fold{}'
    analyze_kfold_results(work_dir_pattern, folds=5)
